import os

import pandas as pd

from .returns import Returns


class ReturnsHelper(object):
    def __init__(self, data_path: str) -> None:
        """initialize

        Args:
            data_path (str): where to load data
        """
        # process [cufe_fm5]
        path = os.path.join(
            data_path, "factor_return", "cufe_fm5", "fivefactor_daily.csv"
        )
        self.cufe_fm5 = pd.read_csv(path)
        self.cufe_fm5.rename(columns={"trddy": "date"}, inplace=True)
        self.cufe_fm5.set_index("date", inplace=True)

        # process [csmar_ind]
        path = os.path.join(data_path, "factor_return", "csmar_ind", "index_daily.pkl")
        self.csmar_ind = pd.read_pickle(path, "gzip")
        self.csmar_ind.set_index("date", inplace=True)
        self.csmar_ind["factor_return"] /= 100

        # process [stock_return]
        path = os.path.join(data_path, "stock_return", "stock_daily.pkl")
        self.stock_return = pd.read_pickle(path, "gzip")
        self.stock_return.set_index("date", inplace=True)

        # process [fund_return]
        path = os.path.join(data_path, "fund_return", "fund_daily.pkl")
        self.fund_return = pd.read_pickle(path, "gzip")
        self.fund_return.sort_values(["fundid", "date"], inplace=True)
        self.fund_return["returns"] = self.fund_return.groupby("fundid")[
            "accnav"
        ].pct_change()
        self.fund_return.dropna(inplace=True)
        self.fund_return.set_index("date", inplace=True)

        return None

    def get_factor(self, source: str, factor_name: str) -> Returns:
        """Get factor return.

        FF: Factor (Fama-French)
        FI: Factor (Index)

        Args:
            source (str): data source of factor
            factor_name (str): factor name

        Returns:
            Returns: factor return
        """
        if source == "cufe_fm5":
            return Returns(f"FF{factor_name}", self.cufe_fm5[factor_name])
        elif source == "csmar_ind":
            data = self.csmar_ind[self.csmar_ind["index_id"] == factor_name]
            return Returns(f"FI{factor_name}", data["factor_return"])

    def get_stock_return(self, stkcd: str) -> Returns:
        """Get stock return

        SK: stock

        Args:
            stkcd (str): stock id

        Returns:
            Returns: stock daily return
        """
        if isinstance(stkcd, str):
            stkcd = int(stkcd)
        data = self.stock_return[self.stock_return["stkcd"] == stkcd]["returns"]
        return Returns(f"SK{stkcd}", data)

    def get_fund_return(self, fundid: str) -> Returns:
        """Get fund return

        FD: Fund

        Args:
            fundid (str): fund id

        Returns:
            Returns: fund return
        """
        if isinstance(fundid, str):
            fundid = int(fundid)
        data = self.fund_return[self.fund_return["fundid"] == fundid]["returns"]
        return Returns(f"FD{fundid}", data)
